import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms
import os
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, accuracy_score, recall_score
import seaborn as sns
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.metrics import roc_curve, auc
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


'''device = 'cuda' if torch.cuda.is_available() else 'cpu'

base_dir = '新图像素材'
image_test_dir = os.path.join(base_dir, 'test')
excel_test = os.path.join(base_dir, 'test.xlsx')
df = pd.read_excel(excel_test)
test_array = df.to_numpy()
x = []
y = test_array[:,13].astype(int)
transform = transforms.Compose([transforms.Resize((256, 256)),  # 统一缩放
                                transforms.ToTensor(),        # 转换为tensor
                                transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                     std=[0.5, 0.5, 0.5])])
for i in range(test_array.shape[0]):
    filename = test_array[i, 1]
    image_path = os.path.join(image_test_dir, filename)
    image = Image.open(image_path).convert('RGB')
    image_tensor = transform(image)
    x.append(image_tensor)

# 加载预训练模型resnet50
model = torchvision.models.resnet50(pretrained=False)
model.fc = nn.Linear(model.fc.in_features, 5)
for param in model.parameters():
    param.requires_grad = False
model.to(device)

PATH='state_dict_model_muti-class.pth'
model.load_state_dict(torch.load(PATH))

model.eval()

predictions = []
with torch.no_grad():
    for i in range(len(x)):
        img = x[i].to(device)
        img = img.unsqueeze(0)
        predict = model(img)
        predict = predict.to('cpu')
        predictions.append(int(np.argmax(predict)))

predictions = np.array(predictions)
print(predictions)

x1 = range(len(x))
plt.figure()
plt.plot(x1, y, marker='+', label='y_true')  # 绘制训练损失线图
plt.plot(x1, predictions, marker='o', color = 'orange', label='y_predict')  # 绘制验证损失散点图
plt.title("predict")  # 设置标题
plt.xlabel("picture number")  # 设置X轴标签
plt.ylabel("点击率")  # 设置Y轴标签
plt.legend()  # 显示图例

conf_matrix = confusion_matrix(y, predictions)
print(conf_matrix)
plt.figure(figsize=(8, 6))
sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', cbar=False)
plt.xlabel('Predicted labels')
plt.ylabel('True labels')
plt.title('Confusion Matrix')
plt.show()

plt.show()'''

#regression
'''device = 'cuda' if torch.cuda.is_available() else 'cpu'

base_dir = '新图像素材'
image_test_dir = os.path.join(base_dir, 'test')
excel_test = os.path.join(base_dir, 'test.xlsx')
df = pd.read_excel(excel_test)
test_array = df.to_numpy()
x = []
y = test_array[:,11].astype(float)
transform_imagenet = transforms.Compose([transforms.Resize((256, 256)),  # 统一缩放
                                transforms.ToTensor(),        # 转换为tensor
                                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                     std=[0.229, 0.224, 0.225])])
transform = transforms.Compose([transforms.Resize((256, 256)),  # 统一缩放
                                transforms.ToTensor(),        # 转换为tensor
                                transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                     std=[0.5, 0.5, 0.5])])
for i in range(test_array.shape[0]):
    filename = test_array[i, 1]
    image_path = os.path.join(image_test_dir, filename)
    image = Image.open(image_path).convert('RGB')
    image_tensor = transform(image)
    x.append(image_tensor)

# 加载预训练模型resnet50
model = torchvision.models.resnet50(pretrained=False)
model.fc = nn.Linear(model.fc.in_features, 1)
for param in model.parameters():
    param.requires_grad = False
model.to(device)

PATH='state_dict_model_regress.pth'
model.load_state_dict(torch.load(PATH))

model.eval()

predictions = []
with torch.no_grad():
    for i in range(len(x)):
        img = x[i].to(device)
        img = img.unsqueeze(0)
        predict = model(img)
        predict = predict.to('cpu')
        predictions.append(predict)

predictions = np.array(predictions)
print(predictions)

x1 = range(len(x))
plt.figure()
plt.plot(x1, y, marker='+', label='y_true')  # 绘制训练损失线图
plt.plot(x1, predictions, marker='o', color = 'orange', label='y_predict')  # 绘制验证损失散点图
plt.title("predict")  # 设置标题
plt.xlabel("picture number")  # 设置X轴标签
plt.ylabel("点击率")  # 设置Y轴标签
plt.legend()  # 显示图例

plt.show()'''

#binary_val

accuracy_loss=np.zeros((20,3))
precision_loss=np.zeros((20,3))
recall_loss=np.zeros((20,3))
f1_loss=np.zeros((20,3))

accuracy_noloss=np.zeros((20,3))
precision_noloss=np.zeros((20,3))
recall_noloss=np.zeros((20,3))
f1_noloss=np.zeros((20,3))

for loop in range(3):
    accuracy_all_loss=[]
    precision_all_loss=[]
    recall_all_loss=[]
    f1_all_loss=[]

    for sjy in range(30,50):
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

        base_dir = '新图像素材'
        image_test_dir = os.path.join(base_dir, 'val')
        excel_test = os.path.join(base_dir, 'val.xlsx')
        df = pd.read_excel(excel_test)
        test_array = df.to_numpy()
        x = []
        t = test_array[:,11].astype(float)
        y = []
        for tt in range(len(t)):
            if t[tt]<0.01*sjy:
                y.append(0)
            else:
                y.append(1)
        y = np.array(y)
        transform = transforms.Compose([transforms.Resize((256, 256)),  # 统一缩放
                                        transforms.ToTensor(),        # 转换为tensor
                                        transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                             std=[0.5, 0.5, 0.5])])
        for i in range(test_array.shape[0]):
            filename = test_array[i, 1]
            image_path = os.path.join(image_test_dir, filename)
            image = Image.open(image_path).convert('RGB')
            image_tensor = transform(image)
            x.append(image_tensor)

        # 加载预训练模型resnet50
        model = torchvision.models.resnet50(pretrained=False)
        model.fc = nn.Sequential(
            nn.Linear(model.fc.in_features, 1),
            nn.Sigmoid()
        )
        for param in model.parameters():
            param.requires_grad = False
        model.to(device)
        if loop == 0:
            PATH=f'loss二分类/state_dict_model_class_{str(0.01*sjy)}.pth'
        else:
            PATH=f'loss二分类/state_dict_model_class_{str(0.01*sjy)}_{str(loop+1)}.pth'
        model.load_state_dict(torch.load(PATH))

        model.eval()

        predictions = []
        prediction_prob = []
        with torch.no_grad():
            for i in range(len(x)):
                img = x[i].to(device)
                img = img.unsqueeze(0)
                predict = model(img)
                predict = predict.to('cpu')
                predict = torch.where(predict.squeeze(1) > 0.5, 1, 0)
                predictions.append(int(predict))

        predictions = np.array(predictions)
        #print(predictions)
        accuracy = accuracy_score(y, predictions)
        precision = precision_score(y, predictions)
        recall = recall_score(y, predictions)
        f1 = f1_score(y, predictions)
        accuracy_all_loss.append(accuracy)
        precision_all_loss.append(precision)
        recall_all_loss.append(recall)
        f1_all_loss.append(f1)
        print(f"Accuracy:{accuracy:.2f}")
        print(f"Precision:{precision:.2f}")
        print(f"Recall:{recall:.2f}")
        print(f"F1 Score:{f1:.2f}")
        conf_matrix = confusion_matrix(y, predictions)
        print(conf_matrix)
        '''plt.figure()
            sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', cbar=False)
            plt.xlabel('Predicted labels')
            plt.ylabel('True labels')
            plt.title('Confusion Matrix')'''

    accuracy_all_noloss=[]
    precision_all_noloss=[]
    recall_all_noloss=[]
    f1_all_noloss=[]
    for sjy in range(30,50):
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

        base_dir = '新图像素材'
        image_test_dir = os.path.join(base_dir, 'val')
        excel_test = os.path.join(base_dir, 'val.xlsx')
        df = pd.read_excel(excel_test)
        test_array = df.to_numpy()
        x = []
        t = test_array[:,11].astype(float)
        y = []
        for tt in range(len(t)):
            if t[tt]<0.01*sjy:
                y.append(0)
            else:
                y.append(1)
        y = np.array(y)
        transform = transforms.Compose([transforms.Resize((256, 256)),  # 统一缩放
                                        transforms.ToTensor(),        # 转换为tensor
                                        transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                             std=[0.5, 0.5, 0.5])])
        for i in range(test_array.shape[0]):
            filename = test_array[i, 1]
            image_path = os.path.join(image_test_dir, filename)
            image = Image.open(image_path).convert('RGB')
            image_tensor = transform(image)
            x.append(image_tensor)

        # 加载预训练模型resnet50
        model = torchvision.models.resnet50(pretrained=False)
        model.fc = nn.Sequential(
            nn.Linear(model.fc.in_features, 1),
            nn.Sigmoid()
        )
        for param in model.parameters():
            param.requires_grad = False
        model.to(device)

        if loop == 0:
            PATH = f'noloss二分类/state_dict_model_class_{str(0.01 * sjy)}.pth'
        else:
            PATH = f'noloss二分类/state_dict_model_class_{str(0.01 * sjy)}_{str(loop+1)}.pth'
        model.load_state_dict(torch.load(PATH))

        model.eval()

        predictions = []
        prediction_prob = []
        with torch.no_grad():
            for i in range(len(x)):
                img = x[i].to(device)
                img = img.unsqueeze(0)
                predict = model(img)
                predict = predict.to('cpu')
                predict = torch.where(predict.squeeze(1) > 0.5, 1, 0)
                predictions.append(int(predict))

        predictions = np.array(predictions)
        #print(predictions)
        accuracy = accuracy_score(y, predictions)
        precision = precision_score(y, predictions)
        recall = recall_score(y, predictions)
        f1 = f1_score(y, predictions)
        accuracy_all_noloss.append(accuracy)
        precision_all_noloss.append(precision)
        recall_all_noloss.append(recall)
        f1_all_noloss.append(f1)
        print(f"Accuracy:{accuracy:.2f}")
        print(f"Precision:{precision:.2f}")
        print(f"Recall:{recall:.2f}")
        print(f"F1 Score:{f1:.2f}")
        conf_matrix = confusion_matrix(y, predictions)
        print(conf_matrix)
        '''plt.figure()
            sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', cbar=False)
            plt.xlabel('Predicted labels')
            plt.ylabel('True labels')
            plt.title('Confusion Matrix')'''
    accuracy_loss[:, loop] = accuracy_all_loss
    precision_loss[:, loop] = precision_all_loss
    recall_loss[:, loop] = recall_all_loss
    f1_loss[:, loop] = f1_all_loss
    accuracy_noloss[:, loop] = accuracy_all_noloss
    precision_noloss[:, loop] = precision_all_noloss
    recall_noloss[:, loop] = recall_all_noloss
    f1_noloss[:, loop] = f1_all_noloss
x = range(20)
plt.figure(1)
mean_y = np.mean(accuracy_loss, axis=1)
std_y = np.std(accuracy_loss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='+', capsize=5, label='Accuracy w/ PL')
mean_y = np.mean(accuracy_noloss, axis=1)
std_y = np.std(accuracy_noloss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='o', capsize=5, label='Accuracy w/o PL')
plt.title("Accuracy")  # 设置标题
plt.xlabel("threshold")  # 设置X轴标签
plt.ylabel("Accuracy")  # 设置Y轴标签
plt.legend()
plt.figure(2)
mean_y = np.mean(precision_loss, axis=1)
std_y = np.std(precision_loss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='+', capsize=5, label='Precision w/ PL')
mean_y = np.mean(precision_noloss, axis=1)
std_y = np.std(precision_noloss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='o', capsize=5, label='Precision w/o PL')
plt.title("Precision")  # 设置标题
plt.xlabel("threshold")  # 设置X轴标签
plt.ylabel("Precision")  # 设置Y轴标签
plt.legend()
plt.figure(3)
mean_y = np.mean(recall_loss, axis=1)
std_y = np.std(recall_loss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='+', capsize=5, label='Recall w/ PL')
mean_y = np.mean(recall_noloss, axis=1)
std_y = np.std(recall_noloss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='o', capsize=5, label='Recall w/o PL')
plt.title("Recall")  # 设置标题
plt.xlabel("threshold")  # 设置X轴标签
plt.ylabel("Recall")  # 设置Y轴标签
plt.legend()
plt.figure(4)
mean_y = np.mean(f1_loss, axis=1)
std_y = np.std(f1_loss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='+', capsize=5, label='F1 w/ PL')
mean_y = np.mean(f1_noloss, axis=1)
std_y = np.std(f1_noloss, axis=1)
plt.errorbar(x, mean_y, yerr=std_y, marker='o', capsize=5, label='F1 w/o PL')
plt.title("F1")  # 设置标题
plt.xlabel("threshold")  # 设置X轴标签
plt.ylabel("F1")  # 设置Y轴标签
plt.legend()
plt.show()